import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import re
# 读取数据
folder_path = "UAV-HAV-5-2-25-5--3model-1"
file_name = "datasource.csv"
df = pd.read_csv(f'{folder_path}/{file_name}')

# 自动识别所有无人机名称
position_cols = [col for col in df.columns if re.match(r'.+_x$', col)]
vehicle_names = list(set([re.sub(r'_x$', '', col) for col in position_cols]))
# 创建 3D 轨迹图
fig3d = plt.figure(figsize=(10, 8))
ax3d = fig3d.add_subplot(111, projection='3d')
# 为每架无人机绘制轨迹
for name in vehicle_names:
    x = df[f"{name}_x"].values
    y = df[f"{name}_y"].values
    z = df[f"{name}_z"].values
    ax3d.plot(x, y, z, label=name)
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('Z')
ax3d.set_title('3D UAV Trajectories')
ax3d.legend()
ax3d.set_aspect('equal')
# 绘制相对信息图（每对无人机）
relative_pairs = [col for col in df.columns if re.match(r'.+_.+_distance$', col)]
if relative_pairs:
    fig_rel, axs_rel = plt.subplots(len(relative_pairs), 1, figsize=(10, 6 * len(relative_pairs)))
    if len(relative_pairs) == 1:
        axs_rel = [axs_rel]  # 确保始终是列表

    for idx, pair_col in enumerate(relative_pairs):
        pair_name = re.sub(r'_distance$', '', pair_col)
        distance_col = f"{pair_name}_distance"
        neighbor_col = f"{pair_name}_neighbor_dist"
        radius_col = f"{pair_name}_radius"

        x = range(len(df))

        axs_rel[idx].plot(x, df[distance_col], label="Distance")
        axs_rel[idx].plot(x, df[neighbor_col], label="Neighbor Distance")
        axs_rel[idx].plot(x, df[radius_col], label="Collision Radius")

        axs_rel[idx].set_title(f"Relative Information: {pair_name}")
        axs_rel[idx].legend()
        axs_rel[idx].grid(True)

    plt.tight_layout()

# 绘制每架无人机的速度图（x, y, z 分量）
fig_vel, axs_vel = plt.subplots(len(vehicle_names), 3, figsize=(12, 4 * len(vehicle_names)))
if len(vehicle_names) == 1:
    axs_vel = [axs_vel]  # 确保始终是列表

for i, name in enumerate(vehicle_names):
    x = range(len(df))
    vx = df[f"{name}_vx"].values
    vy = df[f"{name}_vy"].values
    vz = df[f"{name}_vz"].values

    axs_vel[i][0].plot(x, vx, label="X-Velocity")
    axs_vel[i][0].set_title(f"{name} - X Velocity")
    axs_vel[i][0].legend()
    axs_vel[i][0].grid()

    axs_vel[i][1].plot(x, vy, label="Y-Velocity")
    axs_vel[i][1].set_title(f"{name} - Y Velocity")
    axs_vel[i][1].legend()
    axs_vel[i][1].grid()

    axs_vel[i][2].plot(x, vz, label="Z-Velocity")
    axs_vel[i][2].set_title(f"{name} - Z Velocity")
    axs_vel[i][2].legend()
    axs_vel[i][2].grid()

plt.tight_layout()

# 显示所有图形
plt.show()
# fig1 = plt.figure()
# vehicle_names = ["000", "001"]
# ax = fig1.add_subplot(111, projection='3d')
# for vehicle_name in vehicle_names:
#     ax.plot([json.loads(u)[0] for u in df[vehicle_name].tolist()],
#             [json.loads(u)[1] for u in df[vehicle_name].tolist()],
#             [json.loads(u)[2] for u in df[vehicle_name].tolist()], label=vehicle_name)
#
# ax.legend()
# ax.set_zlim(0, 40)
# # 设置坐标轴标签
# ax.set_xlabel('X')
# ax.set_ylabel('Y')
# ax.set_zlabel('Z')
# ax.set_title('uav_positions')
# ax.set_aspect('equal')
# x = np.arange(0, len(df))
# fig1, axs1 = plt.subplots(4, 1)
# axs1[0].plot(x, df["radius"].tolist(), label="radius")
# axs1[0].plot(x, df["distance"].tolist(), label="distance")
# axs1[0].plot(x, df["neighbor_dist"].tolist(), label="neighbor_dist")
# axs1[0].legend()
# axs1[1].plot(x, [json.loads(u)[0] for u in df["000_velocity"].tolist()], label="000_x_velocity")
# axs1[1].legend()
# axs1[2].plot(x, [json.loads(u)[1] for u in df["000_velocity"].tolist()], label="000_y_velocity")
# axs1[2].legend()
# axs1[3].plot(x, [json.loads(u)[2] for u in df["000_velocity"].tolist()], label="000_z_velocity")
# axs1[3].legend()
#
# fig2, axs2 = plt.subplots(4, 1)
# axs2[0].plot(x, df["radius"].tolist(), label="radius")
# axs2[0].plot(x, df["distance"].tolist(), label="distance")
# axs2[0].plot(x, df["neighbor_dist"].tolist(), label="neighbor_dist")
# axs2[0].legend()
# axs2[1].plot(x, [-json.loads(u)[0] for u in df["001_velocity"].tolist()], label="001_x_velocity")
# axs2[1].legend()
# axs2[2].plot(x, [-json.loads(u)[1] for u in df["001_velocity"].tolist()], label="001_y_velocity")
# axs2[2].legend()
# axs2[3].plot(x, [json.loads(u)[2] for u in df["001_velocity"].tolist()], label="001_z_velocity")
# axs2[3].legend()
# # 调整子图间距
# plt.tight_layout()
# # 显示图形
# plt.show()
